#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import json
import os
from datetime import date
from enum import Enum, IntEnum

import rag.utils
import rag.utils.es_conn
import rag.utils.infinity_conn
import rag.utils.opensearch_coon
from api.constants import RAG_FLOW_SERVICE_NAME
from api.utils import decrypt_database_config, get_base_config
from api.utils.file_utils import get_project_base_directory
from graphrag import search as kg_search
from rag.nlp import search

LIGHTEN = int(os.environ.get("LIGHTEN", "0"))

LLM = None
LLM_FACTORY = None
LLM_BASE_URL = None
CHAT_MDL = ""
EMBEDDING_MDL = ""
RERANK_MDL = ""
ASR_MDL = ""
IMAGE2TEXT_MDL = ""
API_KEY = None
PARSERS = None
HOST_IP = None
HOST_PORT = None
SECRET_KEY = None
FACTORY_LLM_INFOS = None

DATABASE_TYPE = os.getenv("DB_TYPE", "mysql")
DATABASE = decrypt_database_config(name=DATABASE_TYPE)

# authentication
AUTHENTICATION_CONF = None

# client
CLIENT_AUTHENTICATION = None
HTTP_APP_KEY = None
GITHUB_OAUTH = None
FEISHU_OAUTH = None
OAUTH_CONFIG = None
DOC_ENGINE = None
docStoreConn = None

retrievaler = None
kg_retrievaler = None

# user registration switch
REGISTER_ENABLED = 1


# sandbox-executor-manager
SANDBOX_ENABLED = 0
SANDBOX_HOST = None

BUILTIN_EMBEDDING_MODELS = ["BAAI/bge-large-zh-v1.5@BAAI", "maidalun1020/bce-embedding-base_v1@Youdao"]


# 初始化全局设置函数，用于从环境变量和配置文件中加载各种服务和模型相关配置
def init_settings():
    # 声明多个全局变量，以便在函数内部进行赋值操作
    global LLM, LLM_FACTORY, LLM_BASE_URL, LIGHTEN, DATABASE_TYPE, DATABASE, FACTORY_LLM_INFOS, REGISTER_ENABLED
    # 获取LIGHTEN标志，判断是否启用轻量模式，默认为0（不启用）
    LIGHTEN = int(os.environ.get("LIGHTEN", "0"))
    # 从环境变量获取数据库类型，默认为mysql
    DATABASE_TYPE = os.getenv("DB_TYPE", "mysql")
    # 解密并加载指定数据库类型的配置信息
    DATABASE = decrypt_database_config(name=DATABASE_TYPE)
    # 获取用户默认LLM配置字典
    LLM = get_base_config("user_default_llm", {})
    # 从LLM配置中提取默认模型信息
    LLM_DEFAULT_MODELS = LLM.get("default_models", {})
    # 获取LLM工厂名称
    LLM_FACTORY = LLM.get("factory")
    # 获取LLM基础URL
    LLM_BASE_URL = LLM.get("base_url")
    # 尝试从环境变量获取注册是否启用的标志，默认为1（启用）
    try:
        REGISTER_ENABLED = int(os.environ.get("REGISTER_ENABLED", "1"))
    except Exception:
        pass

    # 读取llm_factories.json文件中的工厂模型信息
    try:
        with open(os.path.join(get_project_base_directory(), "conf", "llm_factories.json"), "r") as f:
            FACTORY_LLM_INFOS = json.load(f)["factory_llm_infos"]
    except Exception:
        # 如果读取失败，则使用空列表作为默认值
        FACTORY_LLM_INFOS = []

    # 声明多个全局模型变量，便于后续赋值
    global CHAT_MDL, EMBEDDING_MDL, RERANK_MDL, ASR_MDL, IMAGE2TEXT_MDL
    # 如果未启用轻量模式，则加载内置嵌入模型
    if not LIGHTEN:
        EMBEDDING_MDL = BUILTIN_EMBEDDING_MODELS[0]

    # 如果存在默认模型配置，则从中加载各个模型名称
    if LLM_DEFAULT_MODELS:
        CHAT_MDL = LLM_DEFAULT_MODELS.get("chat_model", CHAT_MDL)
        EMBEDDING_MDL = LLM_DEFAULT_MODELS.get("embedding_model", EMBEDDING_MDL)
        RERANK_MDL = LLM_DEFAULT_MODELS.get("rerank_model", RERANK_MDL)
        ASR_MDL = LLM_DEFAULT_MODELS.get("asr_model", ASR_MDL)
        IMAGE2TEXT_MDL = LLM_DEFAULT_MODELS.get("image2text_model", IMAGE2TEXT_MDL)

        # factory can be specified in the config name with "@". LLM_FACTORY will be used if not specified
        # 如果模型名称没有指定工厂，则附加LLM_FACTORY信息
        CHAT_MDL = CHAT_MDL + (f"@{LLM_FACTORY}" if "@" not in CHAT_MDL and CHAT_MDL != "" else "")
        EMBEDDING_MDL = EMBEDDING_MDL + (f"@{LLM_FACTORY}" if "@" not in EMBEDDING_MDL and EMBEDDING_MDL != "" else "")
        RERANK_MDL = RERANK_MDL + (f"@{LLM_FACTORY}" if "@" not in RERANK_MDL and RERANK_MDL != "" else "")
        ASR_MDL = ASR_MDL + (f"@{LLM_FACTORY}" if "@" not in ASR_MDL and ASR_MDL != "" else "")
        IMAGE2TEXT_MDL = IMAGE2TEXT_MDL + (f"@{LLM_FACTORY}" if "@" not in IMAGE2TEXT_MDL and IMAGE2TEXT_MDL != "" else "")

    # 声明与API、主机、密钥相关的全局变量
    global API_KEY, PARSERS, HOST_IP, HOST_PORT, SECRET_KEY
    # 从LLM配置中获取API密钥
    API_KEY = LLM.get("api_key")
    # 获取解析器配置，默认包含多种文档解析类型
    PARSERS = LLM.get(
        "parsers", "naive:General,qa:Q&A,resume:Resume,manual:Manual,table:Table,paper:Paper,book:Book,laws:Laws,presentation:Presentation,picture:Picture,one:One,audio:Audio,email:Email,tag:Tag"
    )

    # 获取服务主机IP地址，默认为127.0.0.1
    HOST_IP = get_base_config(RAG_FLOW_SERVICE_NAME, {}).get("host", "127.0.0.1")
    # 获取HTTP端口号
    HOST_PORT = get_base_config(RAG_FLOW_SERVICE_NAME, {}).get("http_port")

    # 获取服务密钥，默认为当前日期字符串
    SECRET_KEY = get_base_config(RAG_FLOW_SERVICE_NAME, {}).get("secret_key", str(date.today()))

    # 声明认证相关全局变量
    global AUTHENTICATION_CONF, CLIENT_AUTHENTICATION, HTTP_APP_KEY, GITHUB_OAUTH, FEISHU_OAUTH, OAUTH_CONFIG
    # authentication
    # 获取认证模块的基础配置
    AUTHENTICATION_CONF = get_base_config("authentication", {})

    # client
    # 客户端认证开关及HTTP应用密钥配置
    CLIENT_AUTHENTICATION = AUTHENTICATION_CONF.get("client", {}).get("switch", False)
    HTTP_APP_KEY = AUTHENTICATION_CONF.get("client", {}).get("http_app_key")
    # 获取GitHub OAuth配置
    GITHUB_OAUTH = get_base_config("oauth", {}).get("github")
    # 获取飞书OAuth配置
    FEISHU_OAUTH = get_base_config("oauth", {}).get("feishu")

    # 获取完整的OAuth配置
    OAUTH_CONFIG = get_base_config("oauth", {})

    # 文档引擎相关全局变量声明
    global DOC_ENGINE, docStoreConn, retrievaler, kg_retrievaler
    # 从环境变量获取文档引擎类型，默认为elasticsearch
    DOC_ENGINE = os.environ.get("DOC_ENGINE", "elasticsearch")
    # DOC_ENGINE = os.environ.get('DOC_ENGINE', "opensearch")
    # 转换为小写以统一处理不同写法
    lower_case_doc_engine = DOC_ENGINE.lower()
    # 根据文档引擎类型初始化相应的连接对象
    # 默认用的是elasticsearch
    if lower_case_doc_engine == "elasticsearch":
        docStoreConn = rag.utils.es_conn.ESConnection()
    # 也可以使用infinity
    elif lower_case_doc_engine == "infinity":
        docStoreConn = rag.utils.infinity_conn.InfinityConnection()
    elif lower_case_doc_engine == "opensearch":
        docStoreConn = rag.utils.opensearch_coon.OSConnection()
    else:
        # 不支持的文档引擎类型抛出异常
        raise Exception(f"Not supported doc engine: {DOC_ENGINE}")

    # 初始化检索器和知识图谱检索器实例
    retrievaler = search.Dealer(docStoreConn)
    kg_retrievaler = kg_search.KGSearch(docStoreConn)

    # 如果启用了沙箱功能
    if int(os.environ.get("SANDBOX_ENABLED", "0")):
        # 声明沙箱主机全局变量
        global SANDBOX_HOST
        # 设置沙箱主机地址，默认为sandbox-executor-manager
        SANDBOX_HOST = os.environ.get("SANDBOX_HOST", "sandbox-executor-manager")


# 定义一个自定义的枚举类，扩展Enum的功能，提供验证和获取枚举值及名称的方法
class CustomEnum(Enum):
    # 验证给定的值是否为当前枚举的有效成员
    @classmethod
    def valid(cls, value):
        # 尝试使用给定值创建枚举成员，如果成功则表示值有效
        try:
            cls(value)
            return True
        except BaseException:
            # 如果转换过程中发生任何异常，则表示值无效
            return False

    # 返回所有枚举成员的值列表
    @classmethod
    def values(cls):
        # 利用列表推导式遍历所有枚举成员并收集它们的值
        return [member.value for member in cls.__members__.values()]

    # 返回所有枚举成员的名称列表
    @classmethod
    def names(cls):
        # 利用列表推导式遍历所有枚举成员并收集它们的名称
        return [member.name for member in cls.__members__.values()]


class RetCode(IntEnum, CustomEnum):
    SUCCESS = 0
    NOT_EFFECTIVE = 10
    EXCEPTION_ERROR = 100
    ARGUMENT_ERROR = 101
    DATA_ERROR = 102
    OPERATING_ERROR = 103
    CONNECTION_ERROR = 105
    RUNNING = 106
    PERMISSION_ERROR = 108
    AUTHENTICATION_ERROR = 109
    UNAUTHORIZED = 401
    SERVER_ERROR = 500
    FORBIDDEN = 403
    NOT_FOUND = 404
